import torch
import numpy as np
from torch.utils.data import Dataset
from typing import Dict, Tuple, Union, Optional, Callable, Any
import random
from . import DATASETS


@DATASETS.register()
class SyntheticDataset(Dataset):
    """
    自定义合成数据集类，用于生成伪造的图像和标签数据
    
    参数:
        image_shape: 图像的形状，例如 (C, H, W) 或 (C, H, W, D)
        num_samples: 数据集中的样本数量
        num_classes: 分类数量（用于标签生成）
        transform: 应用于数据的变换
        key_names: 字典中使用的键名，默认为 ["image", "label"]
    """
    
    def __init__(
        self,
        image_shape: Tuple[int, ...] = (1, 256, 256),
        num_samples: int = 1000,
        num_classes: int = 1,
        transform: Optional[Callable] = None,
        key_names: Tuple[str, str] = ("image", "label")
    ):
        self.image_shape = image_shape
        self.num_samples = num_samples
        self.num_classes = num_classes
        self.transform = transform
        self.key_names = key_names
        
        # 验证参数
        if num_samples <= 0:
            raise ValueError("num_samples must be positive")
        if num_classes <= 0:
            raise ValueError("num_classes must be positive")
        if len(image_shape) < 2:
            raise ValueError("image_shape must have at least 2 dimensions")
            
    def __len__(self) -> int:
        """返回数据集大小"""
        return self.num_samples
    
    def _generate_synthetic_image(self) -> np.ndarray:
        """生成合成图像数据"""
        # 生成随机图像数据
        image = np.random.rand(*self.image_shape).astype(np.float32)
        return image
    
    def _generate_synthetic_label(self) -> np.ndarray:
        """生成合成标签数据"""
        # 对于分类任务，生成分类标签
        if self.num_classes == 1:
            # 二分类情况，生成0-1标签
            label = (np.random.rand(*self.image_shape) > 0.5).astype(np.float32)
        else:
            # 多分类情况，生成整数标签，范围[0, num_classes-1]
            label_shape = self.image_shape
            label = np.random.randint(0, self.num_classes, size=label_shape).astype(np.int64)
        return label
    
    def __getitem__(self, index: int) -> Dict[str, Union[torch.Tensor, np.ndarray]]:
        """
        获取指定索引的数据样本
        
        参数:
            index: 样本索引
            
        返回:
            包含图像和标签的字典
        """
        # 生成合成数据
        image = self._generate_synthetic_image()
        label = self._generate_synthetic_label()
        
        # 构造样本字典
        sample = {
            self.key_names[0]: image,
            self.key_names[1]: label,
            # 添加元数据以支持MONAI变换
            f"{self.key_names[0]}_meta_dict": {},
            f"{self.key_names[1]}_meta_dict": {}
        }
        
        # 应用变换
        if self.transform is not None:
            sample = self.transform(sample)
            
        return sample


# 兼容性别名
SyntheticDataset = SyntheticDataset